package org.xq.softcup.mq.registry.impl;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xq.softcup.mq.registry.ServiceRegistry;
import org.xq.softcup.mq.util.XQRpcException;
import org.xq.softcup.mq.util.ZkClient;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;

/**
 * 基于zookeeper注册服务, 用于分布式消息队列
 */
@Slf4j
public class ZkServiceRegistry extends ServiceRegistry {
    public static final String ENV = "env";
    public static final String ZK_ADDRESS = "zkaddress";
    public static final String ZK_DIGEST = "zkdigest";
    private static final String zkBasePath = "/xqmq-rpc";
    private String zkEnvPath;
    private ZkClient zkClient = null;
    private Thread refreshThread;
    private volatile boolean refreshThreadStop = false;
    private volatile ConcurrentMap<String, TreeSet<String>> registryData = new ConcurrentHashMap();
    private volatile ConcurrentMap<String, TreeSet<String>> discoveryData = new ConcurrentHashMap();

    public String keyToPath(String nodeKey) {
        return this.zkEnvPath + "/" + nodeKey;
    }

    public String pathToKey(String nodePath) {
        if ((nodePath == null) || (nodePath.length() <= this.zkEnvPath.length()) || (!nodePath.startsWith(this.zkEnvPath))) {
            return null;
        }
        return nodePath.substring(this.zkEnvPath.length() + 1, nodePath.length());
    }

    public void start(Map<String, String> param) {
        String zkaddress = (String) param.get("zkaddress");
        String zkdigest = (String) param.get("zkdigest");
        String env = (String) param.get("env");
        if ((zkaddress == null) || (zkaddress.trim().length() == 0)) {
            throw new XQRpcException("zkaddress不能为空!");
        }
        if ((env == null) || (env.trim().length() == 0)) {
            throw new XQRpcException("env不能为空!");
        }
        this.zkEnvPath = "/xqmq-rpc".concat("/").concat(env);


        this.zkClient = new ZkClient(zkaddress, this.zkEnvPath, zkdigest, new Watcher() {
            public void process(WatchedEvent watchedEvent) {
                try {
                    ZkServiceRegistry.log.debug("观察者:{}", watchedEvent);
                    if (watchedEvent.getState() == Watcher.Event.KeeperState.Expired) {
                        ZkServiceRegistry.this.zkClient.destroy();
                        ZkServiceRegistry.this.zkClient.getClient();

                        ZkServiceRegistry.this.refreshDiscoveryData(null);
                        ZkServiceRegistry.log.info("Zookeeper 重连加载成功!");
                    }
                    String path = watchedEvent.getPath();
                    String key = ZkServiceRegistry.this.pathToKey(path);
                    if (key != null) {
                        ZkServiceRegistry.this.zkClient.getClient().exists(path, true);
                        if (watchedEvent.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
                            ZkServiceRegistry.this.refreshDiscoveryData(key);
                        } else if (watchedEvent.getState() == Watcher.Event.KeeperState.SyncConnected) {
                            ZkServiceRegistry.log.info("重新加载");
                        }
                    }
                } catch (Exception e) {
                    ZkServiceRegistry.log.error(e.getMessage(), e);
                }
            }
        });
        this.zkClient.getClient();


        this.refreshThread = new Thread(new Runnable() {
            public void run() {
                while (!ZkServiceRegistry.this.refreshThreadStop) {
                    try {
                        TimeUnit.SECONDS.sleep(60L);
                        ZkServiceRegistry.this.refreshDiscoveryData(null);
                        ZkServiceRegistry.this.refreshRegistryData();
                    } catch (Exception e) {
                        if (!ZkServiceRegistry.this.refreshThreadStop) {
                            ZkServiceRegistry.log.error("刷新线程错误.", e);
                        }
                    }
                }
                ZkServiceRegistry.log.info("刷新线程停止.");
            }
        });
        this.refreshThread.setName("ZkServiceRegistry 刷新线程.");
        this.refreshThread.setDaemon(true);
        this.refreshThread.start();

        log.info("ZkServiceRegistry 初始化完成. [env={}]", env);
    }

    public void stop() {
        if (this.zkClient != null) {
            this.zkClient.destroy();
        }
        if (this.refreshThread != null) {
            this.refreshThreadStop = true;
            this.refreshThread.interrupt();
        }
    }

    private void refreshDiscoveryData(String key) {
        Set<String> keys = new HashSet();
        if ((key != null) && (key.trim().length() > 0)) {
            keys.add(key);
        } else if (this.discoveryData.size() > 0) {
            keys.addAll(this.discoveryData.keySet());
        }
        if (keys.size() > 0) {
            for (String keyItem : keys) {
                String path = keyToPath(keyItem);
                Map<String, String> childPathData = this.zkClient.getChildPathData(path);


                TreeSet<String> existValues = (TreeSet) this.discoveryData.get(keyItem);
                if (existValues == null) {
                    existValues = new TreeSet();
                    this.discoveryData.put(keyItem, existValues);
                }
                if (childPathData.size() > 0) {
                    existValues.clear();
                    existValues.addAll(childPathData.keySet());
                }
            }
            log.info("刷新检索数据, discoveryData = {}", this.discoveryData);
        }
    }

    private void refreshRegistryData() {
        if (this.registryData.size() > 0) {
            for (Map.Entry<String, TreeSet<String>> item : this.registryData.entrySet()) {
                String key = (String) item.getKey();
                TreeSet<String> valueSet = item.getValue();
                for (String value : valueSet) {
                    String path = keyToPath(key);
                    this.zkClient.setChildPathData(path, value, "");
                }
            }
            String key;
            log.info("刷新注册数据, registryData = {}", this.registryData);
        }
    }

    public boolean registry(Set<String> keys, String value) {
        for (String key : keys) {
            TreeSet<String> values = (TreeSet) this.registryData.get(key);
            if (values == null) {
                values = new TreeSet();
                this.registryData.put(key, values);
            }
            values.add(value);


            String path = keyToPath(key);
            this.zkClient.setChildPathData(path, value, "");
        }
        log.info("注册成功, keys = {}, value = {}", keys, value);
        return true;
    }

    public boolean remove(Set<String> keys, String value) {
        for (String key : keys) {
            TreeSet<String> values = (TreeSet) this.discoveryData.get(key);
            if (values != null) {
                values.remove(value);
            }
            String path = keyToPath(key);
            this.zkClient.deleteChildPath(path, value);
        }
        log.info("删除成功, keys = {}, value = {}", keys, value);
        return true;
    }

    public Map<String, TreeSet<String>> discovery(Set<String> keys) {
        if ((keys == null) || (keys.size() == 0)) {
            return null;
        }
        Map<String, TreeSet<String>> registryDataTmp = new HashMap();
        for (String key : keys) {
            TreeSet<String> valueSetTmp = discovery(key);
            if (valueSetTmp != null) {
                registryDataTmp.put(key, valueSetTmp);
            }
        }
        return registryDataTmp;
    }

    public TreeSet<String> discovery(String key) {
        TreeSet<String> values = (TreeSet) this.discoveryData.get(key);
        if (values == null) {
            refreshDiscoveryData(key);

            values = (TreeSet) this.discoveryData.get(key);
        }
        return values;
    }
}
